TopkFusion ================= 在指定维度上选取 Top-K 个最大值,并同时输出对应的索引。 该算子在融合实现中同时完成: - Top-K 值计算 - 索引提取 - 可选的结果排序 其计算过程为: 在指定维度上,对每个 slice 进行排序或选择,输出前 ``K`` 个最大元素及其原始索引。 当 ``sorted = false`` 时,仅保证输出为 Top-K 元素集合, 不保证其在输出中的顺序;当 ``sorted = true`` 时, 输出结果按值从大到小排序。 输入: - **input** - 输入张量的数据地址。 数据类型需与所调用的 TopkFusion 接口类型一致,例如:``fp_*`` 接口对应 ``float*``,``hp_*`` 接口对应 ``half*``,``i16_*`` 接口对应 ``int16_t*``,``i32_*`` 接口对应 ``int32_t*`` - **output** - Top-K 值输出张量的数据地址。 数据类型与 ``input`` 保持一致。 - **output_index** - Top-K 索引输出张量的数据地址,类型固定为 ``int32_t*``。 - **parameter** - Top-K 参数结构体指针 ``TopkParameter*``,其定义如下: .. code-block:: c typedef struct TopkParameter { // primitive parameter OpParameter op_parameter_; int k_; // 需要选取的 Top-K 个数 int axis_; // 执行 Top-K 的维度 bool sorted_; // 是否对 Top-K 结果按值排序 // other parameter int dim_size_; // axis 维度长度 int outer_loop_num_; // axis 之前维度展开后的循环次数 int inner_loop_num_; // axis 之后维度展开后的循环次数 void *topk_node_list_; // 临时 TopK 节点缓冲区 } TopkParameter; - **core_mask** - 核掩码(仅共享存储版本使用)。 输出: - **output** - Top-K 值输出张量的数据地址。 - **output_index** - Top-K 索引输出张量的数据地址。 支持平台: ``FT78NE`` ``MT7004`` .. note:: - FT78NE 支持的数据类型: fp32, fp64, int8, int16, int32 - MT7004 支持的数据类型: fp16, fp32, int16, int32 - ``output`` 与 ``output_index`` 的布局与输入张量保持一致,仅在 Top-K 维度长度变为 ``K`` - TopkFusion 算子内部使用 ``void*`` 进行类型复用, 但调用时 ``input`` 与 ``output`` 的实际数据类型 必须与所选接口前缀严格一致。 **共享存储版本:** .. c:function:: void fp_topk_fusion_s(void* input, void* output, int32_t* output_index, TopkParameter* parameter, int core_mask) .. c:function:: void hp_topk_fusion_s(void* input, void* output, int32_t* output_index, TopkParameter* parameter, int core_mask) .. c:function:: void i16_topk_fusion_s(void* input, void* output, int32_t* output_index, TopkParameter* parameter, int core_mask) .. c:function:: void i32_topk_fusion_s(void* input, void* output, int32_t* output_index, TopkParameter* parameter, int core_mask) **C调用示例:** .. code-block:: c :linenos: :emphasize-lines: 19 // FT78NE 多核示例 #include #include int main(int argc, char* argv[]) { float *input = (float *)0xA0000000; float *output = (float *)0xB0000000; int32_t *output_index = (int32_t *)0xB1000000; TopkParameter param; param.k_ = 5; param.dim_size_ = 100; param.outer_loop_num_ = 1; param.inner_loop_num_ = 1; param.sorted_ = 1; int core_mask = 0xff; fp_topk_fusion_s(input, output, output_index, ¶m, core_mask); return 0; } **私有存储版本:** .. c:function:: void fp_topk_fusion_p(void* input, void* output, int32_t* output_index, TopkParameter* parameter) .. c:function:: void hp_topk_fusion_p(void* input, void* output, int32_t* output_index, TopkParameter* parameter) .. c:function:: void i16_topk_fusion_p(void* input, void* output, int32_t* output_index, TopkParameter* parameter) .. c:function:: void i32_topk_fusion_p(void* input, void* output, int32_t* output_index, TopkParameter* parameter) **C调用示例:** .. code-block:: c :linenos: :emphasize-lines: 17 // MT7004 单核示例 #include #include int main(int argc, char* argv[]) { half *input = (half *)0x10000000; half *output = (half *)0x10010000; int32_t *output_index = (int32_t *)0x10020000; TopkParameter param; param.k_ = 3; param.dim_size_ = 64; param.outer_loop_num_ = 1; param.inner_loop_num_ = 1; param.sorted_ = 0; hp_topk_fusion_p(input, output, output_index, ¶m); return 0; }